// SPDX-License-Identifier: MIT

#include <math.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/ec.h>
#include <openssl/rsa.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/err.h>
#include <openssl/provider.h>
#include <oqs/oqs.h>

#if defined(OQS_USE_RASPBERRY_PI)
#define _OQS_RASPBERRY_PI
#endif
#if defined(OQS_SPEED_USE_ARM_PMU)
#define SPEED_USE_ARM_PMU
#endif
#include "ds_benchmark.h"
#include "system_info.c"
#include "test_common.h"

struct SigPair {
    unsigned char *sigma1;
    size_t sigma1_len;
    unsigned char *sigma2;
    size_t sigma2_len;
} signature;

const char *kHybridSignatureAlgorithms[] = {
    "p256_dilithium2", "rsa3072_dilithium2", "p384_dilithium3", "p521_dilithium5",
    "p256_mldsa44", "rsa3072_mldsa44", "p384_mldsa65", "p521_mldsa87",
    "p256_falcon512", "rsa3072_falcon512", "p256_falconpadded512", "rsa3072_falconpadded512",
    "p521_falcon1024", "p521_falconpadded1024", "p256_sphincssha2128fsimple", "rsa3072_sphincssha2128fsimple",
    "p256_sphincssha2128ssimple", "rsa3072_sphincssha2128ssimple", "p384_sphincssha2192fsimple",
    "p256_sphincsshake128fsimple", "rsa3072_sphincsshake128fsimple", NULL
};

static OSSL_LIB_CTX *init_openssl(void) {
    OSSL_LIB_CTX *ctx;

    if (!(ctx = OSSL_LIB_CTX_new())) {
        fprintf(stderr, cRED "failed to initialize a new `OSSL_LIB_CTX` " cNORM "\n");
    }
    return ctx;
}


static OQS_STATUS sig_speed_wrapper(const char *method_name, uint64_t duration, bool printInfo, bool doFullCycle) {

	OQS_SIG *sig = NULL;
	uint8_t *public_key = NULL;
	uint8_t *secret_key = NULL;
	uint8_t *message = NULL;
	uint8_t *signature = NULL;
	size_t message_len = 50;
	size_t signature_len = 0;
	OQS_STATUS ret = OQS_ERROR;

   sig = OQS_SIG_new(method_name);
	if (sig == NULL) {
		return OQS_SUCCESS;
	}

	public_key = malloc(sig->length_public_key);
	secret_key = malloc(sig->length_secret_key);
	message = malloc(message_len);
	signature = malloc(sig->length_signature);

	if ((public_key == NULL) || (secret_key == NULL) || (message == NULL) || (signature == NULL)) {
		fprintf(stderr, "ERROR: malloc failed\n");
		goto err;
	}

	OQS_randombytes(message, message_len);

	printf("%-36s | %10s | %14s | %15s | %10s | %25s | %10s\n", sig->method_name, "", "", "", "", "", "");
	if (!doFullCycle) {
		TIME_OPERATION_SECONDS(OQS_SIG_keypair(sig, public_key, secret_key), "keypair", duration)
		TIME_OPERATION_SECONDS(OQS_SIG_sign(sig, signature, &signature_len, message, message_len, secret_key), "sign", duration)
		TIME_OPERATION_SECONDS(OQS_SIG_verify(sig, message, message_len, signature, signature_len, public_key), "verify", duration)
	} else {

	}


	if (printInfo) {
		printf("public key bytes: %zu, secret key bytes: %zu, signature bytes: %zu\n", sig->length_public_key, sig->length_secret_key, sig->length_signature);
		if (signature_len != sig->length_signature) {
			printf("   Actual signature length returned (%zu) less than declared maximum signature length (%zu)\n", signature_len, sig->length_signature);
		}
	}

	ret = OQS_SUCCESS;
	goto cleanup;

err:
	ret = OQS_ERROR;

cleanup:
	if (sig != NULL) {
		OQS_MEM_secure_free(secret_key, sig->length_secret_key);
	}
	OQS_MEM_insecure_free(public_key);
	OQS_MEM_insecure_free(signature);
	OQS_MEM_insecure_free(message);
	OQS_SIG_free(sig);

	return ret;
}

static OQS_STATUS printAlgs(void) {
	for (size_t i = 0; i < OQS_SIG_algs_length; i++) {
		OQS_SIG *sig = OQS_SIG_new(OQS_SIG_alg_identifier(i));
		if (sig == NULL) {
			printf("%s (disabled)\n", OQS_SIG_alg_identifier(i));
		} else {
			printf("%s\n", OQS_SIG_alg_identifier(i));
		}
		OQS_SIG_free(sig);
	}
	return OQS_SUCCESS;
}

static OQS_STATUS verify_hybrid_signature(OSSL_LIB_CTX *libctx, EVP_PKEY *pkey, const uint8_t *message, size_t message_len){
    EVP_MD_CTX *mdctx = EVP_MD_CTX_new(); 
    //fprintf(stderr, "%d \n", signature.sigma1_len);
    EVP_DigestVerifyInit_ex(mdctx, NULL, "SHA256", libctx, NULL, pkey, NULL) ;
    EVP_DigestVerifyUpdate(mdctx, message, message_len);
         if(EVP_DigestVerifyFinal(mdctx, signature.sigma1, signature.sigma1_len)!=1){
        fprintf(stderr, "Classic pkey verify failed.\n");
        EVP_MD_CTX_free(mdctx);
                     return OQS_ERROR;
    }
    EVP_MD_CTX_free(mdctx);
          return OQS_SUCCESS;
}

static OQS_STATUS generate_hybrid_signature(OSSL_LIB_CTX *libctx, EVP_PKEY *pkey, const uint8_t *message, size_t message_len){
    EVP_MD_CTX *mdctx = EVP_MD_CTX_new(); 
    EVP_DigestSignInit_ex(mdctx, NULL, "SHA256", libctx, NULL, pkey, NULL) ;
    EVP_DigestSignUpdate(mdctx, message, message_len);
    EVP_DigestSignFinal(mdctx, NULL, &signature.sigma1_len);
    signature.sigma1 = OPENSSL_malloc(signature.sigma1_len);
           if(signature.sigma1 == NULL){
         fprintf(stderr, "Pkey signed failed.\n");
         EVP_MD_CTX_free(mdctx);
                        return OQS_ERROR;
    }
    EVP_DigestSignFinal(mdctx, signature.sigma1, &signature.sigma1_len);
    EVP_MD_CTX_free(mdctx); 
          return OQS_SUCCESS;          
}

static EVP_PKEY *generate_key(OSSL_LIB_CTX *libctx, const char *signame){
    EVP_PKEY_CTX *ctx = NULL;
    EVP_PKEY *pqkey = NULL;
 
   if (!(ctx = EVP_PKEY_CTX_new_from_name(libctx, signame, NULL)) ||!EVP_PKEY_keygen_init(ctx) ||!(EVP_PKEY_generate(ctx, &pqkey))) {
      fprintf(stderr, "EVP_PKEY generate error.\n");
      EVP_PKEY_CTX_free(ctx);
      EVP_PKEY_free(pqkey);
               return NULL;
    }

   EVP_PKEY_CTX_free(ctx);
        return pqkey;    
}

static OQS_STATUS test_algorithm(OSSL_LIB_CTX *libctx,const char *signames, uint64_t duration, bool printInfo, bool doFullCycle){

	OQS_SIG *sig = NULL;
	uint8_t *message = NULL;
	size_t message_len = 50;
	OQS_STATUS ret = OQS_ERROR;	

   signature.sigma1 = NULL;
   signature.sigma1_len = 0;
   signature.sigma2 = NULL;
   signature.sigma2_len = 0;
   
   EVP_PKEY *pkey = NULL;
   message = malloc(message_len);
	OQS_randombytes(message, message_len);

	printf("|%-36s | %10s | %14s | %15s | %10s | %25s | %10s|\n", signames, "", "", "", "", "", "");
       if (!doFullCycle) {
		TIME_OPERATION_SECONDS((pkey = generate_key(libctx, signames)) != NULL, "keypair", duration)
		TIME_OPERATION_SECONDS(generate_hybrid_signature(libctx, pkey, message, message_len), "sign", duration)
		TIME_OPERATION_SECONDS(verify_hybrid_signature(libctx, pkey, message, message_len), "verify", duration)
	}
	 
	 if (printInfo) {
		 printf("|signature bytes:%-21s| %10zu | %14s | %15s | %10s | %25s | %10s|\n", "",  signature.sigma1_len ,"", "", "", "", "");
	}
	 
	OQS_MEM_insecure_free(message);		
	OQS_SIG_free(sig);
	return OQS_SUCCESS;
err:
	OQS_MEM_insecure_free(message);		
	OQS_SIG_free(sig);
	return OQS_ERROR;
}


int main(int argc, char **argv) {

	int ret = EXIT_SUCCESS;
	OSSL_LIB_CTX *libctx;
   OSSL_PROVIDER *oqs_provider;
	const char *modulename;
        const char *configfile;
	bool printUsage = false;
	uint64_t duration = 5;
	bool printSigInfo = true;
	bool doFullCycle = false;

       if (!(libctx = init_openssl())) {
                   goto end;
   }	
	T(argc == 3);
   modulename = argv[1];
   configfile = argv[2];
	
        if(!(oqs_provider = OSSL_PROVIDER_load(libctx, modulename))){
        fputs(cRED "  `oqs_provider` is NULL " cNORM "\n", stderr);
                      goto end;    
    }	
	
   //sigalgs = OSSL_PROVIDER_query_operation(oqs_provider, OSSL_OP_SIGNATURE, &query_nocache);		
		
   print_system_info();

	printf("Speed test\n");
	printf("==========\n");

	PRINT_TIMER_HEADER	
       for (int i = 0 ; kHybridSignatureAlgorithms[i] != NULL; i++) {
                 if (test_algorithm(libctx, kHybridSignatureAlgorithms[i], duration, printSigInfo, doFullCycle)) {
            fprintf(stderr, cRED " Failed for %s " cNORM "\n", kHybridSignatureAlgorithms[i]);
       } else {
          //  fprintf(stderr, cGREEN "%s succeeded" cNORM "\n", kHybridSignatureAlgorithms[i]);       
       }
   } 	

	 PRINT_TIMER_FOOTER
end:
    OSSL_PROVIDER_unload(oqs_provider);   
    OSSL_LIB_CTX_free(libctx);
	return ret;
}
